package com.lee.base.sharding;

import cn.hutool.db.Db;
import cn.hutool.db.ds.DSFactory;
import com.google.common.collect.Range;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import javax.sql.DataSource;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.config.sharding.ShardingRuleConfiguration;
import org.apache.shardingsphere.api.config.sharding.TableRuleConfiguration;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;
import org.apache.shardingsphere.shardingjdbc.api.ShardingDataSourceFactory;
import org.junit.Before;
import org.junit.Test;

/**
 * @author : libo@juzishuke.com
 * @Title :
 * @Description :
 * @date : 2024/4/7 14:26
 * @since :1.0
 */
@Slf4j
public class SingleKeyComplexAlgorithm {

    DataSource dataSource = null;

    @Before
    public void setUp() throws SQLException {
        DataSource ds = DSFactory.get();
        Map<String, DataSource> dataSourceMap = new HashMap<>();
        dataSourceMap.put("ds0", ds);
        TableRuleConfiguration orderTableRuleConfig = new TableRuleConfiguration("t_order",
                "ds0.t_order_${0..3}");

        // 配置分库 + 分表策略
//        orderTableRuleConfig.setTableShardingStrategyConfig(
//                new ComplexShardingStrategyConfiguration("order_id",new HintShardingAlgorithm<Integer>(){}));

        // 配置分片规则
        ShardingRuleConfiguration shardingRuleConfig = new ShardingRuleConfiguration();
        shardingRuleConfig.getTableRuleConfigs().add(orderTableRuleConfig);
        Properties props = new Properties();
        props.put("sql.show", true);
        dataSource = ShardingDataSourceFactory.createDataSource(dataSourceMap, shardingRuleConfig,
                props);

    }
    @Test
    public void testUpdate1() throws SQLException {
        Db.use(dataSource).execute("update  t_order set id=2 where order_id =? ",2);
    }
    @Test
    public void testUpdate2() throws SQLException {
        Db.use(dataSource).execute("update  t_order set order_id=2 where order_id =? ",2);
    }
    @Test
    public void testEq() throws SQLException {
        Db.use(dataSource).query("select * from t_order where id = ? and order_id=? ",1,22);
        Db.use(dataSource).query("select * from t_order where id = ? and order_id=? ",1,23);
    }
    @Test
    public void testNotAll() throws SQLException {
        Db.use(dataSource).query("select * from t_order where id = ? ",1);
        Db.use(dataSource).query("select * from t_order where  order_id=? ",23);
    }
    @Test
    public void testIn() throws SQLException {
        Db.use(dataSource).query("select * from t_order where id in (?) and order_id in (?,?) ",1,22,23);
        Db.use(dataSource).query("select * from t_order where id = ? and order_id in (?,?) ",1,22,23);
    }

    @Test
    public void testRange() throws SQLException {
        Db.use(dataSource).query("select * from t_order where id = ? and  order_id between ? and ? ",1,1,911);
        Db.use(dataSource).query("select * from t_order where id in ( ?)  and  order_id between ? and ? ",1,1,911);
    }

    class OrderIdRangeAlgorithm implements ComplexKeysShardingAlgorithm<Integer> {

        @Override
        public Collection<String> doSharding(Collection<String> availableTargetNames,
                ComplexKeysShardingValue<Integer> shardingValue) {
            log.info("分片键参数:{}", shardingValue.toString());
            Collection<String> result = new ArrayList<>();
            Collection<Integer> ids = shardingValue.getColumnNameAndShardingValuesMap()
                    .getOrDefault("id", new ArrayList<>());
            Collection<Integer> orderIds = shardingValue.getColumnNameAndShardingValuesMap()
                    .getOrDefault("order_id", new ArrayList<>());
            Range<Integer> orderIdRange = shardingValue.getColumnNameAndRangeValuesMap()
                    .get("order_id");
            if (ids.isEmpty() || (orderIds.isEmpty() && orderIdRange == null)) {
                return availableTargetNames;
            }
            ;
            for (Integer userId : ids) {
                for (Integer orderId : orderIds) {
                    extracted(availableTargetNames, userId, orderId, result);
                }
                if (orderIdRange != null && !orderIdRange.isEmpty()) {
                    if (orderIdRange.hasLowerBound()) {
                        extracted(availableTargetNames, userId, orderIdRange.lowerEndpoint(),
                                result);
                    }
                    if (orderIdRange.hasUpperBound()) {
                        extracted(availableTargetNames, userId, orderIdRange.upperEndpoint(),
                                result);
                    }
                }
            }
            return result;
        }

        private void extracted(Collection<String> availableTargetNames, Integer userId,
                Integer orderId, Collection<String> result) {
            String suffix;
            if (userId % 2 == 1 && orderId % 2 == 1) {
                // 奇数
                suffix = "1";
            } else if (userId % 2 == 0 && orderId % 2 == 0) {
                // 偶数
                suffix = "2";
            } else {
                // 1奇数1偶数
                suffix = "0";
            }
            for (String targetName : availableTargetNames) {
                if (targetName.endsWith("_" + suffix)) { // 按照分片值选择目标表
                    result.add(targetName);
                }
            }
        }
    }
}
